/**
 * This program is free software, you can redistribute it and/or modify it.
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
 * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/*!
 * \file power.h
 * \brief
*/

#include "kernel_operator.h"

constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue
using namespace AscendC;


template <typename TYPE_X, typename TYPE_Y, typename TYPE_Z>
class KernelPowSafe {
public:
    __aicore__ inline KernelPowSafe() {}

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z,
                                const TilingData &tiling_data
                                ) 
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        uint32_t coreNum = GetBlockIdx();
        this->globalBufferIndex =tiling_data.bigCoreDataNum * coreNum;
        this->tileDataNum = tiling_data.tileDataNum;

        this->is_input0_scalar = tiling_data.is_input0_scalar;
        this->is_input1_scalar =  tiling_data.is_input1_scalar;
        this->yDim  = tiling_data.yDim;
        this->isSameX1  = tiling_data.isSameX1;
        this->isSameX2  = tiling_data.isSameX2;

        const uint32_t* strideX1Src =tiling_data.strideX1;
        const uint32_t* strideX2Src =tiling_data.strideX2;
        const uint32_t* strideYSrc  =tiling_data.strideY;
        for (int i = 0; i < 10; ++i) {
            this->strideX1[i] = strideX1Src[i];
            this->strideX2[i] = strideX2Src[i];
            this->strideY[i]  = strideYSrc[i];
        }


        if (coreNum < tiling_data.tailBlockNum) {
            this->coreDataNum = tiling_data.bigCoreDataNum;
            this->tileNum = tiling_data.finalBigTileNum;
            this->tailDataNum = tiling_data.bigTailDataNum;
        }
        else {
            this->coreDataNum = tiling_data.smallCoreDataNum;
            this->tileNum = tiling_data.finalSmallTileNum;
            this->tailDataNum = tiling_data.smallTailDataNum;
            globalBufferIndex -= (tiling_data.bigCoreDataNum - tiling_data.smallCoreDataNum) * (GetBlockIdx() - tiling_data.tailBlockNum);
        }

        if (tiling_data.isSameX1) {
            // 无广播，按输出偏移访问
            xGm.SetGlobalBuffer((__gm__ TYPE_X*)x + globalBufferIndex, coreDataNum);
        } else {
            // 有广播，直接访问原始数据（不叠加输出偏移）
            xGm.SetGlobalBuffer((__gm__ TYPE_X*)x, tiling_data.X1TotalNum);  // x1TotalNum是x的原始总元素数
        }
        if (tiling_data.isSameX2) {
            yGm.SetGlobalBuffer((__gm__ TYPE_Y*)y + globalBufferIndex, coreDataNum);
        } else {
            yGm.SetGlobalBuffer((__gm__ TYPE_Y*)y, tiling_data.X2TotalNum);  // x2TotalNum是y的原始总元素数
        }

        zGm.SetGlobalBuffer((__gm__ TYPE_Z *)z + globalBufferIndex, this->coreDataNum);

        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileDataNum * sizeof(TYPE_X));
        pipe.InitBuffer(inQueueY, BUFFER_NUM, this->tileDataNum * sizeof(TYPE_Y));
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileDataNum * sizeof(TYPE_Z));

        pipe.InitBuffer(absTmp, this->tileDataNum * sizeof(float));
        pipe.InitBuffer(zeroTmp, this->tileDataNum * sizeof(float));
        pipe.InitBuffer(maskTmp, this->tileDataNum * sizeof(uint8_t));

        pipe.InitBuffer(xCastTmp, this->tileDataNum * sizeof(float));
        pipe.InitBuffer(yCastTmp, this->tileDataNum * sizeof(float));
        pipe.InitBuffer(zCastTmp, this->tileDataNum * sizeof(float));
        pipe.InitBuffer(tranCastTmp, this->tileDataNum * sizeof(half));

    }


    __aicore__ inline void Process() {
        int32_t loopCount = this->tileNum;
        this->processDataNum = this->tileDataNum;
        for (int32_t i = 0; i < loopCount-1; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
        this->processDataNum = this->tailDataNum;
        CopyIn(loopCount-1);
        Compute(loopCount-1);
        CopyOut(loopCount-1);
    }

private:
    // ===============================================
    //  广播索引计算函数
    //  将输出索引映射到输入索引
    // ===============================================
    __aicore__ inline uint32_t GetBroadcastIndexEff( uint32_t linearIdx,
                                                const uint32_t *effStride  // 预计算后的有效stride
                                                )
    {
        // 利用预计算stride快速定位
        // linearIdx 按输出 shape 展开
        uint32_t offset = 0;
        uint32_t tmp = linearIdx;
        for (int i = 0; i < yDim; ++i) {
            uint32_t stride_y = this->strideY[i];
            if (stride_y == 0) {
                continue;
            }
            uint32_t coord = tmp / stride_y;
            tmp %= stride_y;
            offset += coord * effStride[i];
        }
        return offset;
    }

    __aicore__ inline void CopyIn(int32_t progress) {
        LocalTensor<TYPE_X> xLocal = inQueueX.AllocTensor<TYPE_X>();
        LocalTensor<TYPE_Y> yLocal = inQueueY.AllocTensor<TYPE_Y>();
        if (this->is_input0_scalar && !this->is_input1_scalar) {
            TYPE_X scalarValX;
            scalarValX = xGm.GetValue(0);//duplicate不支持int8和uint8
            for (int i = 0; i < this->processDataNum; ++i) {
                xLocal.SetValue(i, scalarValX);
            }
            DataCopy(yLocal, yGm[progress * this->tileDataNum], this->processDataNum);
        }
        else if (!this->is_input0_scalar && this->is_input1_scalar) {
            DataCopy(xLocal, xGm[progress * this->tileDataNum], this->processDataNum);
            TYPE_Y scalarValY;
            scalarValY = yGm.GetValue(0);
             for (int i = 0; i < this->processDataNum; ++i) {
                yLocal.SetValue(i, scalarValY);
            }
        }
        else{
            // ====== 判断是否完全匹配 ======
            if (isSameX1 && isSameX2 ) {
                // 无广播，直接DMA搬运
                DataCopy(xLocal, xGm[progress * tileDataNum], processDataNum);
                DataCopy(yLocal, yGm[progress * tileDataNum], processDataNum);
            } else if ( !isSameX1 && isSameX2 ){//x广播yDMA搬运
                 for (uint32_t i = 0; i < processDataNum; ++i) {
                    uint32_t globalIdx = globalBufferIndex + i;
                    uint32_t idxX = GetBroadcastIndexEff(globalIdx, this->strideX1);
                    xLocal.SetValue(i, xGm.GetValue(idxX));
                }
                DataCopy(yLocal, yGm[progress * tileDataNum], processDataNum);  
            }else if ( isSameX1 && !isSameX2 ){//xDMA搬运y广播
                DataCopy(xLocal, xGm[progress * tileDataNum], processDataNum);
                 for (uint32_t i = 0; i < processDataNum; ++i) {
                    uint32_t globalIdx = globalBufferIndex + i;
                    uint32_t idxY = GetBroadcastIndexEff(globalIdx, this->strideX2);
                    yLocal.SetValue(i, yGm.GetValue(idxY));
                }
            }else {//全广播
                 for (uint32_t i = 0; i < processDataNum; ++i) {
                    uint32_t globalIdx = globalBufferIndex + i;
                    uint32_t idxX = GetBroadcastIndexEff(globalIdx, this->strideX1);
                    uint32_t idxY = GetBroadcastIndexEff(globalIdx, this->strideX2);
                    xLocal.SetValue(i, xGm.GetValue(idxX));
                    yLocal.SetValue(i, yGm.GetValue(idxY));
                    
                }
            }          
        }
        
        inQueueX.EnQue(xLocal);
        inQueueY.EnQue(yLocal);
    }

    __aicore__ inline void CopyOut(uint32_t progress) {

        LocalTensor<TYPE_Z> zLocal = outQueueZ.DeQue<TYPE_Z>();
        DataCopy(zGm[progress * this->tileDataNum], zLocal, this->processDataNum);
        outQueueZ.FreeTensor(zLocal);
    }

    __aicore__ inline void PowCompute(LocalTensor<float> &xLocal, LocalTensor<float> &yLocal,
                                                            LocalTensor<float> &zLocal, LocalTensor<uint8_t> &mask)
    {

        LocalTensor<float> p1 = absTmp.Get<float>();
        LocalTensor<float> zeros = zeroTmp.Get<float>();
        LocalTensor<float> &scalars = zeros;
        LocalTensor<float> &ones = yLocal;
        Duplicate(scalars, 2.0f, this->processDataNum);
        Abs(p1, xLocal, this->processDataNum);
        Ln(p1, p1, this->processDataNum);
        Mul(p1, p1, yLocal, this->processDataNum);
        Exp(zLocal, p1, this->processDataNum); // zLocal = exp(y*len(|x|))

        Abs(p1, yLocal, this->processDataNum);
        
        Div(zeros, p1, scalars, this->processDataNum);
        Cast(zeros, zeros, RoundMode::CAST_TRUNC, this->processDataNum);
        Mul(zeros, zeros, scalars, this->processDataNum);
        Sub(p1, p1, zeros, this->processDataNum);
        
        Muls(p1, p1, -2.0f, this->processDataNum);
        Adds(p1, p1, 1.0f, this->processDataNum);
        Mul(p1, p1, zLocal, this->processDataNum);

        Duplicate(zeros, 0.0f, this->processDataNum);

        Compare(mask, xLocal, zeros, CMPMODE::LT, this->processDataNum); // x小于0的部分
        Select(zLocal, mask, p1, zLocal, SELMODE::VSEL_TENSOR_TENSOR_MODE, this->processDataNum);

        Compare(mask, yLocal, zeros, CMPMODE::EQ, this->processDataNum); // exp等于0的部分
        Duplicate(ones, 1.0f, this->processDataNum);
        Select(p1, mask, ones, zeros, SELMODE::VSEL_TENSOR_TENSOR_MODE, this->processDataNum);

        Compare(mask, xLocal, zeros, CMPMODE::EQ, this->processDataNum); // x等于0的部分
        Select(zLocal, mask, p1, zLocal, SELMODE::VSEL_TENSOR_TENSOR_MODE, this->processDataNum);
    }



    __aicore__ inline void Compute(int32_t progress)
    {
        LocalTensor<TYPE_X> xLocal = inQueueX.DeQue<TYPE_X>();
        LocalTensor<TYPE_Y> yLocal = inQueueY.DeQue<TYPE_Y>();
        LocalTensor<TYPE_Z> zLocal = outQueueZ.AllocTensor<TYPE_Z>();

        LocalTensor<float> xLocalFp = xCastTmp.Get<float>();
        LocalTensor<float> yLocalFp = yCastTmp.Get<float>();
        LocalTensor<float> zLocalFp = zCastTmp.Get<float>();

        LocalTensor<uint8_t> mask = maskTmp.Get<uint8_t>();

        for(int i = 0 ; i < 10 ; i++){
            PRINTF("[DEBUG]xLocal=%d, yLocal=%d \n", xLocal.GetValue(i), yLocal.GetValue(i));
        }

        if constexpr ( IsSameType<TYPE_X, half>::value || IsSameType<TYPE_X, bfloat16_t>::value || 
                IsSameType<TYPE_X, int16_t>::value ||IsSameType<TYPE_X, int32_t>::value) {
            Cast(xLocalFp, xLocal, RoundMode::CAST_NONE, this->processDataNum); 
        } 
        if constexpr ( IsSameType<TYPE_X, uint8_t>::value ||IsSameType<TYPE_X, int8_t>::value) {
            LocalTensor<half> tmpLocal = tranCastTmp.Get<half>();
            Cast(tmpLocal, xLocal, RoundMode::CAST_NONE, this->processDataNum);
            Cast(xLocalFp, tmpLocal, RoundMode::CAST_NONE, this->processDataNum); 
        } 
        if constexpr ( IsSameType<TYPE_X, float32_t>::value) {
            xLocalFp = xLocal.template ReinterpretCast<float>(); 
        }

        if constexpr ( IsSameType<TYPE_Y, half>::value || IsSameType<TYPE_Y, bfloat16_t>::value ||
                IsSameType<TYPE_Y, int16_t>::value ||IsSameType<TYPE_Y, int32_t>::value) {
            Cast(yLocalFp, yLocal, RoundMode::CAST_NONE, this->processDataNum); 
        }
        if constexpr ( IsSameType<TYPE_Y, uint8_t>::value ||IsSameType<TYPE_Y, int8_t>::value) {
            LocalTensor<half> tmpLocal = tranCastTmp.Get<half>(); 
            Cast(tmpLocal, yLocal, RoundMode::CAST_NONE, this->processDataNum); 
            Cast(yLocalFp, tmpLocal, RoundMode::CAST_NONE, this->processDataNum); 
        } 
        if constexpr ( IsSameType<TYPE_Y, float32_t>::value) {
            yLocalFp = yLocal.template ReinterpretCast<float>(); 
        }

        PowCompute(xLocalFp, yLocalFp, zLocalFp, mask);

        if constexpr ( IsSameType<TYPE_Z, half>::value || IsSameType<TYPE_Z, bfloat16_t>::value || 
                IsSameType<TYPE_Z, int16_t>::value ||IsSameType<TYPE_Z, int32_t>::value) {
            Cast(zLocal, zLocalFp, RoundMode::CAST_RINT, this->processDataNum); 
        } 
        if constexpr ( IsSameType<TYPE_Z, uint8_t>::value ||IsSameType<TYPE_Z, int8_t>::value) {
            LocalTensor<half> tmpLocal = tranCastTmp.Get<half>(); 
            Cast(tmpLocal, zLocalFp, RoundMode::CAST_NONE, this->processDataNum); 
            Cast(zLocal, tmpLocal, RoundMode::CAST_RINT, this->processDataNum); 
        } 
        if constexpr ( IsSameType<TYPE_Z, float32_t>::value) {
            zLocal = zLocalFp.template ReinterpretCast<float>(); 
        }

        outQueueZ.EnQue<TYPE_Z>(zLocal);
        inQueueX.FreeTensor(xLocal); 
        inQueueY.FreeTensor(yLocal);
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX,inQueueY;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueZ;
    TBuf<QuePosition::VECCALC> absTmp;
    TBuf<QuePosition::VECCALC> xCastTmp;
    TBuf<QuePosition::VECCALC> yCastTmp;
    TBuf<QuePosition::VECCALC> zCastTmp;
    TBuf<QuePosition::VECCALC> tranCastTmp;
    TBuf<QuePosition::VECCALC> maskTmp;
    TBuf<QuePosition::VECCALC> zeroTmp;
    GlobalTensor<TYPE_X> xGm;
    GlobalTensor<TYPE_Y> yGm;
    GlobalTensor<TYPE_Z> zGm;
    uint32_t coreDataNum;
    uint32_t tileNum;
    uint32_t tileDataNum;
    uint32_t tailDataNum;
    uint32_t processDataNum;

    uint64_t globalBufferIndex;
    uint32_t is_input0_scalar;
    uint32_t is_input1_scalar;
    //形状广播
    uint32_t yDim;
    uint32_t isSameX1;
    uint32_t isSameX2;
    uint32_t strideX1[10];
    uint32_t strideX2[10];
    uint32_t strideY[10];
};
            


                
// =========================================================
//                 Kernel Entry
// =========================================================
extern "C" __global__ __aicore__ void pow_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);

    const int x_dtype = tiling_data.x_dtype;
    const int y_dtype = tiling_data.y_dtype;
    PRINTF("[DEBUG]x_dtype=%d, y_dtype=%d \n", x_dtype, y_dtype);

    // ==== 类型分派 ====
#define CASE_POW(DT_X, TYPE_X, DT_Y, TYPE_Y, TYPE_Z) \
    if (x_dtype == DT_X && y_dtype == DT_Y) { \
        PRINTF("[KERNEL]:Enter pow case: X=%s, Y=%s, Z=%s\n",#TYPE_X, #TYPE_Y, #TYPE_Z);\
        KernelPowSafe<TYPE_X, TYPE_Y, TYPE_Z> op; \
        op.Init(x, y, z, tiling_data); \
        op.Process(); \
        return; \
    }
    // ================== 类型映射规则 ==================
    // ====== INT8 行 ======
    CASE_POW(DT_INT8, int8_t,  DT_INT8,   int8_t,    int8_t);
    CASE_POW(DT_INT8, int8_t,  DT_UINT8,  uint8_t,   int16_t);
    CASE_POW(DT_INT8, int8_t,  DT_INT16,  int16_t,   int16_t);
    CASE_POW(DT_INT8, int8_t,  DT_INT32,  int32_t,   int32_t);
    CASE_POW(DT_INT8, int8_t,  DT_FLOAT16, half,      half);
    CASE_POW(DT_INT8, int8_t,  DT_BF16,   bfloat16_t, bfloat16_t);
    CASE_POW(DT_INT8, int8_t,  DT_FLOAT,  float,      float);

    // ====== UINT8 行 ======
    CASE_POW(DT_UINT8, uint8_t, DT_INT8,   int8_t,    int16_t);
    CASE_POW(DT_UINT8, uint8_t, DT_UINT8,  uint8_t,   uint8_t);
    CASE_POW(DT_UINT8, uint8_t, DT_INT16,  int16_t,   int16_t);
    CASE_POW(DT_UINT8, uint8_t, DT_INT32,  int32_t,   int32_t);
    CASE_POW(DT_UINT8, uint8_t, DT_FLOAT16, half,      half);
    CASE_POW(DT_UINT8, uint8_t, DT_BF16,   bfloat16_t, bfloat16_t);
    CASE_POW(DT_UINT8, uint8_t, DT_FLOAT,  float,      float);

    // ====== INT16 行 ======
    CASE_POW(DT_INT16, int16_t, DT_INT8,   int8_t,    int16_t);
    CASE_POW(DT_INT16, int16_t, DT_UINT8,  uint8_t,   int16_t);
    CASE_POW(DT_INT16, int16_t, DT_INT16,  int16_t,   int16_t);
    CASE_POW(DT_INT16, int16_t, DT_INT32,  int32_t,   int32_t);
    CASE_POW(DT_INT16, int16_t, DT_FLOAT16, half,      half);
    CASE_POW(DT_INT16, int16_t, DT_BF16,   bfloat16_t, bfloat16_t);
    CASE_POW(DT_INT16, int16_t, DT_FLOAT,  float,      float);

    // ====== INT32 行 ======
    CASE_POW(DT_INT32, int32_t, DT_INT8,   int8_t,    int32_t);
    CASE_POW(DT_INT32, int32_t, DT_UINT8,  uint8_t,   int32_t);
    CASE_POW(DT_INT32, int32_t, DT_INT16,  int16_t,   int32_t);
    CASE_POW(DT_INT32, int32_t, DT_INT32,  int32_t,   int32_t);
    CASE_POW(DT_INT32, int32_t, DT_FLOAT16, half,      half);
    CASE_POW(DT_INT32, int32_t, DT_BF16,   bfloat16_t, bfloat16_t);
    CASE_POW(DT_INT32, int32_t, DT_FLOAT,  float,      float);

    // ====== FLOAT16 行 ======
    CASE_POW(DT_FLOAT16, half, DT_INT8,   int8_t,    half);
    CASE_POW(DT_FLOAT16, half, DT_UINT8,  uint8_t,   half);
    CASE_POW(DT_FLOAT16, half, DT_INT16,  int16_t,   half);
    CASE_POW(DT_FLOAT16, half, DT_INT32,  int32_t,   half);
    CASE_POW(DT_FLOAT16, half, DT_FLOAT16, half,      half);
    CASE_POW(DT_FLOAT16, half, DT_BF16,   bfloat16_t, float);
    CASE_POW(DT_FLOAT16, half, DT_FLOAT,  float,      float);

    // ====== BF16 行 ======
    CASE_POW(DT_BF16, bfloat16_t, DT_INT8,   int8_t,    bfloat16_t);
    CASE_POW(DT_BF16, bfloat16_t, DT_UINT8,  uint8_t,   bfloat16_t);
    CASE_POW(DT_BF16, bfloat16_t, DT_INT16,  int16_t,   bfloat16_t);
    CASE_POW(DT_BF16, bfloat16_t, DT_INT32,  int32_t,   bfloat16_t);
    CASE_POW(DT_BF16, bfloat16_t, DT_FLOAT16, half,      float);
    CASE_POW(DT_BF16, bfloat16_t, DT_BF16,   bfloat16_t, bfloat16_t);
    CASE_POW(DT_BF16, bfloat16_t, DT_FLOAT,  float,      float);

    // ====== FLOAT 行 ======
    CASE_POW(DT_FLOAT, float, DT_INT8,   int8_t,    float);
    CASE_POW(DT_FLOAT, float, DT_UINT8,  uint8_t,   float);
    CASE_POW(DT_FLOAT, float, DT_INT16,  int16_t,   float);
    CASE_POW(DT_FLOAT, float, DT_INT32,  int32_t,   float);
    CASE_POW(DT_FLOAT, float, DT_FLOAT16, half,      float);
    CASE_POW(DT_FLOAT, float, DT_BF16,   bfloat16_t, float);
    CASE_POW(DT_FLOAT, float, DT_FLOAT,  float,      float);
    // ================== 结束 ==================
#undef CASE_POW
     ASSERT(false && "Unsupported dtype combination in pow_custom!");
}
